#! /usr/bin/env python3
# coding: utf-8

import numpy as np
import sys, os
from chainer import cuda
from progressbar import progressbar
import librosa
import soundfile as sf
import pickle as pic

from FCA_online import FCA_online
from configure import *


class MNMF_online(FCA_online):
    """ Blind Source Separation Using Multichannel Nonnegative Matrix Factorization (MNMF)

    X_FTM: the observed complex spectrogram
    covarianceMatrix_NFMM: spatial covariance matrices (SCMs) for each source
    W_NFK: basis vectors for each source
    H_NKT: activations for each source
    lambda_NFT: power spectral densities of each source (W_NFK @ H_NKT)
    """

    def __init__(self, n_source=2, n_basis=2, xp=np, init_SCM="unit", mode_update_parameter=["all", "one_by_one"][0], total_frame=0, all_sep_spec=0, first_batch_size=7, weight_rho=0.9, mini_batch_size=4):
        """ initialize MNMF

        Parameters:
        -----------
            n_source: int
                the number of sources
            n_basis: int
                the number of bases of each source
            xp : numpy or cupy
            init_SCM: str
                how to initialize covariance matrix {unit, obs, ILRMA}
            mode_update_parameter: str
                'all' : update all the parameters simultanesouly to reduce computational cost
                'one_by_one' : update the parameters one by one to monotonically increase log-likelihood
        """
        super(MNMF_online, self).__init__(n_source=n_source, xp=xp, init_SCM=init_SCM, mode_update_parameter=mode_update_parameter)
        self.n_basis = n_basis
        self.method_name = "MNMF_online"
        self.total_frame = total_frame
        self.all_sep_spec = all_sep_spec
        self.first_batch_size = first_batch_size
        self.weight_rho = weight_rho
        self.mini_batch_size = mini_batch_size

    def set_parameter(self, n_source=None, n_iteration=None, n_basis=None, init_SCM=None, mode_update_parameter=None):
        """ set parameters

        Parameters:
        -----------
            n_source: int
                the number of sources
            init_SCM: str
                how to initialize covariance matrix {unit, obs, ILRMA}
            mode_update_parameter: str
                'all' : update all the variables simultanesouly
                'one_by_one' : update one by one
        """
        if n_source != None:
            self.n_source = n_source
        if n_iteration != None:
            self.n_iteration = n_iteration
        if n_basis != None:
            self.n_basis = n_basis
        if init_SCM != None:
            self.init_SCM = init_SCM
        if mode_update_parameter != None:
            self.mode_update_parameter = mode_update_parameter


    ## 初始化first_batch的pre_batch参数，均设为0，等效为离线
    def initialize_pre_batch(self):
        # G参数初始化
        self.pre_batch_phi = self.xp.zeros([self.n_source, self.n_freq, self.n_mic, self.n_mic])
        self.pre_batch_psi = self.xp.zeros([self.n_source, self.n_freq, self.n_mic, self.n_mic])
        self.pre_batch_covarianceMatrix_NFMM = self.xp.zeros([self.n_source, self.n_freq, self.n_mic, self.n_mic])
        # W参数初始化
        self.pre_batch_W_a_1 = self.xp.zeros([self.n_source, self.n_freq, self.n_basis])  # N F K
        self.pre_batch_W_b_1 = self.xp.zeros([self.n_source, self.n_freq, self.n_basis])  # N F K
        self.pre_batch_W_NFK = self.xp.zeros([self.n_source, self.n_freq, self.n_basis])  # N F K


    def initialize_PSD(self):
        power_observation_FT = (self.xp.abs(self.X_FTM).astype(self.xp.float) ** 2).mean(axis=2)
        shape = 2
        self.W_NFK = self.xp.random.dirichlet(np.ones(self.n_freq) * shape, size=[self.n_source, self.n_basis]).transpose(0, 2, 1)
        self.H_NKT = self.xp.random.gamma(shape, (power_observation_FT.mean() * self.n_freq * self.n_mic / (self.n_source * self.n_basis)) / shape, size=[self.n_source, self.n_basis, self.n_time])
        self.H_NKT[self.H_NKT < EPS] = EPS

        self.lambda_NFT = self.W_NFK @ self.H_NKT

    def make_filename_suffix(self):
        self.filename_suffix = "S={}-it={}-K={}-init={}-update={}-lb={}".format(self.n_source, self.n_iteration, self.n_basis, self.init_SCM, self.mode_update_parameter, self.mini_batch_size)

        if hasattr(self, "file_id"):
            file_id_str = str(self.file_id)
            self.filename_suffix += "-ID={}".format(file_id_str)
        else:
            print("====================\n\nWarning: Please set self.file_id\n\n====================")

        print("filename_suffix:", self.filename_suffix)

    def update(self):
        if self.mode_update_parameter == "one_by_one":
            self.update_axiliary_variable()
            self.update_W()
            self.update_axiliary_variable()
            self.update_H()
            self.update_axiliary_variable()
            self.update_covarianceMatrix()
        if self.mode_update_parameter == "all":
            self.update_axiliary_variable()
            self.update_WH()
            self.update_covarianceMatrix()
        self.normalize()

    def update_WH(self):
        self.W_a_1 = (self.H_NKT.transpose(0, 2, 1)[:, None] * self.tr_Cov_Yinv_X_Yinv_NFT[:, :, :, None]).sum(axis=2)  # N F K
        self.W_b_1 = (self.H_NKT.transpose(0, 2, 1)[:, None] * self.tr_Cov_Yinv_NFT[:, :, :, None]).sum(axis=2)  # N F K
        # online version
        Fj_WaW_NFK = self.W_NFK * self.W_a_1 * self.W_NFK + self.weight_rho * self.pre_batch_W_NFK * self.pre_batch_W_a_1 * self.pre_batch_W_NFK  # (43) of [2]
        Fj_b_NFK = self.W_b_1 + self.weight_rho * self.pre_batch_W_b_1  # (43) of [2]

        a_2 = (self.W_NFK[..., None] * self.tr_Cov_Yinv_X_Yinv_NFT[:, :, None]).sum(axis=1)  # N K T
        b_2 = (self.W_NFK[..., None] * self.tr_Cov_Yinv_NFT[:, :, None]).sum(axis=1)  # N K T

        self.W_NFK = self.xp.sqrt(Fj_WaW_NFK / Fj_b_NFK)
        self.H_NKT = self.H_NKT * self.xp.sqrt(a_2 / b_2)

    def update_W(self):
        self.W_a_1 = (self.H_NKT.transpose(0, 2, 1)[:, None] * self.tr_Cov_Yinv_X_Yinv_NFT[:, :, :, None]).sum(axis=2)  # N F K
        self.W_b_1 = (self.H_NKT.transpose(0, 2, 1)[:, None] * self.tr_Cov_Yinv_NFT[:, :, :, None]).sum(axis=2)  # N F K
        # online version
        Fj_WaW_NFK = self.W_NFK * self.W_a_1 * self.W_NFK + self.weight_rho * self.pre_batch_W_NFK * self.pre_batch_W_a_1 * self.pre_batch_W_NFK  # (43) of [2]
        Fj_b_NFK = self.W_b_1 + self.weight_rho * self.pre_batch_W_b_1  # (43) of [2]
        self.W_NFK = self.xp.sqrt(Fj_WaW_NFK / Fj_b_NFK)
        self.lambda_NFT = self.W_NFK @ self.H_NKT + EPS

    def update_H(self):
        a_1 = (self.W_NFK[..., None] * self.tr_Cov_Yinv_X_Yinv_NFT[:, :, None]).sum(axis=1)  # N K T
        b_1 = (self.W_NFK[..., None] * self.tr_Cov_Yinv_NFT[:, :, None]).sum(axis=1)  # N K T
        self.H_NKT = self.H_NKT * self.xp.sqrt(a_1 / b_1)
        self.lambda_NFT = self.W_NFK @ self.H_NKT + EPS

    def normalize(self):
        mu_NF = self.xp.trace(self.covarianceMatrix_NFMM, axis1=2, axis2=3).real
        self.covarianceMatrix_NFMM = self.covarianceMatrix_NFMM / mu_NF[:, :, None, None]
        self.W_NFK = self.W_NFK * mu_NF[:, :, None]

        nu_NK = self.W_NFK.sum(axis=1)
        self.W_NFK = self.W_NFK / nu_NK[:, None]
        self.H_NKT = self.H_NKT * nu_NK[:, :, None]

        self.lambda_NFT = self.W_NFK @ self.H_NKT + EPS

    def save_parameter(self, filename):
        param_list = [self.covarianceMatrix_NFMM, self.lambda_NFT]
        param_list.append(self.W_NFK)
        param_list.append(self.H_NKT)

        if self.xp != np:
            param_list = [cuda.to_cpu(param) for param in param_list]
        pic.dump(param_list, open(filename, "wb"))

    def load_parameter(self, filename):
        param_list = pic.load(open(filename, "rb"))
        if self.xp != np:
            param_list = [cuda.to_gpu(param) for param in param_list]

        self.covarianceMatrix_NFMM, self.lambda_NFT, self.W_NFK, self.H_NKT = param_list


class pre_batch_init():
    def __init__(self, W_a_1=None, W_b_1=None, W_NFK=None, phi=None, psi=None, covarianceMatrix_NFMM=None):
        # W参数的传递
        self.W_a_1 = W_a_1
        self.W_b_1 = W_b_1
        self.W_NFK = W_NFK
        # G参数的传递
        self.phi = phi
        self.psi = psi
        self.covarianceMatrix_NFMM = covarianceMatrix_NFMM


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument(         'input_fileName', type=str, help='filename of the multichannel observed signals')
    parser.add_argument(              '--file_id', type=str, default="None", help='file id')
    parser.add_argument(                  '--gpu', type=int, default=    -1, help='GPU ID')
    parser.add_argument(                '--n_fft', type=int, default=  1024, help='number of frequencies')
    parser.add_argument(              '--n_basis', type=int, default=    16, help='number of basis of NMF')
    parser.add_argument(             '--n_source', type=int, default=     2, help='number of noise')
    parser.add_argument(          '--n_iteration', type=int, default=30, help='number of iteration')
    parser.add_argument('--mode_update_parameter', type=str, default="all", help='all, one_by_one')
    parser.add_argument(             '--init_SCM', type=str, default="obs", help='unit, obs, ILRMA')
    parser.add_argument(            '--save_path', type=str, default="..\\output\\", help='save path of wave file')
    parser.add_argument(             '--save_wav', type=str, default=True, help='save wave file')
    parser.add_argument(     '--first_batch_size', type=int, default=40, help='frame num of first mini batch')
    parser.add_argument(      '--mini_batch_size', type=int, default=4, help='frame num of each mini batch')
    parser.add_argument(           '--weight_rho', type=float, default=0.9, help='online update weight rho of last mini batch')
    args = parser.parse_args()

    if args.gpu < 0:
        import numpy as xp
    else:
        import cupy as xp

        print("Use GPU " + str(args.gpu))
        cuda.get_device_from_id(args.gpu).use()

    wav, fs = sf.read(args.input_fileName)
    wav = wav.T
    M = len(wav)
    for m in range(M):
        tmp = librosa.core.stft(wav[m], n_fft=args.n_fft, hop_length=int(args.n_fft / 4))
        if m == 0:
            spec_all = np.zeros([tmp.shape[0], tmp.shape[1], M], dtype=np.complex)
        spec_all[:, :, m] = tmp

    T = spec_all.shape[1]
    first_batch_size = args.first_batch_size
    mini_batch_size = args.mini_batch_size
    all_sep_spec = 0

    for t in range(first_batch_size,T+1):
        # online分帧
        if t == first_batch_size:
            tmp = spec_all[:,:t,:]
        else:
            tmp = spec_all[:,t-mini_batch_size:t,:]
        spec = np.zeros([tmp.shape[0], tmp.shape[1], tmp.shape[2]], dtype=np.complex)
        spec = tmp
        args.file_id = t+1-first_batch_size

        # 初始化separater对象
        separater = MNMF_online(n_source=args.n_source, n_basis=args.n_basis, xp=xp, init_SCM=args.init_SCM, mode_update_parameter=args.mode_update_parameter, total_frame=T, all_sep_spec=all_sep_spec, first_batch_size=first_batch_size, weight_rho=args.weight_rho, mini_batch_size=args.mini_batch_size)

        # 初始化STFT频谱(继承自父类FCA_online),包括x_ft和X_ft = x_ft * x_ft^H in [1]
        separater.load_spectrogram(spec)

        # 定义命名相关的参数
        separater.file_id = args.file_id
        separater.fs = fs

        # 定义pre_batch对象，用于存储上一batch的参数值
        if args.file_id == 1:
            pre_batch = pre_batch_init()
        else:
            # W参数的传递
            separater.pre_batch_W_a_1 = pre_batch.W_a_1
            separater.pre_batch_W_b_1 = pre_batch.W_b_1
            separater.pre_batch_W_NFK = pre_batch.W_NFK
            # G参数的传递
            separater.pre_batch_phi = pre_batch.phi
            separater.pre_batch_psi = pre_batch.psi
            separater.pre_batch_covarianceMatrix_NFMM = pre_batch.covarianceMatrix_NFMM

        # solve函数,负责算法核心运算、参数保存和结果输出(继承自父类FCA_online)
        all_sep_spec = separater.solve(n_iteration=args.n_iteration, save_likelihood=False, save_parameter=False, save_wav=args.save_wav, save_path=args.save_path, interval_save_parameter=100)

        # 在线的参数保存，提供给下一个mini-batch
        # W参数的保存
        pre_batch.W_a_1 = separater.W_a_1
        pre_batch.W_b_1 = separater.W_b_1
        pre_batch.W_NFK = separater.W_NFK
        # G参数的保存
        pre_batch.phi = separater.phi
        pre_batch.psi = separater.psi
        pre_batch.covarianceMatrix_NFMM = separater.covarianceMatrix_NFMM

